from fastapi import Depends, UploadFile
from overrides import override
import typing as T
from datetime import datetime
import uuid
from milvus_model.hybrid.mgte import MGTEEmbeddingFunction


from ...mapper.UserKnowledgeBaseMapper import (
    _UserKnowledgeBaseMapper,
    UserKnowledgeBaseMapper,
)
from ....pojo.dto import (
    UserKnowledgeBaseDTO,
    KnowledgeHeaderDTO,
    EmbeddingSearchParamsDTO,
    UpdateImageInfoDTO,
)
from ....pojo.entity import UserKnowledgeBase, UserKnowledge, UserKnowledgeBaseFile
from ....pojo.vo import ImageInfoVO
from ..UserKnowledgeBaseService import (
    UserKnowledgeBaseService as _UserKnowledgeBaseService,
)
from ....common.utils import loader_spliter
from ....init.Global import EmbeddingFunction


class UserKnowledgeBaseServiceImpl(_UserKnowledgeBaseService):
    mapper: _UserKnowledgeBaseMapper = None
    embeddingFunction: MGTEEmbeddingFunction = None

    def __call__(
        self, mapper: UserKnowledgeBaseMapper, embeddingFunction: EmbeddingFunction
    ):
        self.embeddingFunction = embeddingFunction
        self.mapper = mapper
        return self

    @override
    async def createKnowledgeBase(self, userKnowledgeBaseDTO: UserKnowledgeBaseDTO):
        createTime = datetime.now()
        alias = userKnowledgeBaseDTO.alias or f"{uuid.uuid4()}"
        entity = UserKnowledgeBase(
            alias=alias, createTime=createTime, userId=userKnowledgeBaseDTO.userId
        )
        await self.mapper.insertByEntity(entity)

    @override
    async def getKnowledgeBaseListByUserId(self, userId: int):
        return await self.mapper.getKnowledgeBaseListByUserId(userId)

    @override
    async def updateKnowledgeBaseNameByKnowledgeBaseId(
        self, knowledgeBaseId: int, name: str
    ):
        await self.mapper.updateKnowledgeBaseNameByKnowledgeBaseId(
            knowledgeBaseId, name
        )

    @override
    async def deleteKnowledgeBaseByKnowledgeId(self, knowledgeBaseId: int):
        await self.mapper.deleteKnowledgeBaseById(knowledgeBaseId)

    @override
    async def upload(self, file: UploadFile, knowledgeHeader: KnowledgeHeaderDTO):
        # TODO 优化表结构，现在查询太麻烦
        file_name = file.filename
        type = file.content_type.split("/")[1]
        _uuid = f"{uuid.uuid4()}"
        if type in ["png", "jpg", "jpeg", "jpeg"]:
            object_name = f"{_uuid}.{type}"
            docs = loader_spliter(
                await file.read(), _uuid, file_name, type, object_name=object_name
            )
        else:
            docs = loader_spliter(await file.read(), _uuid, file_name, type)

        contents = [doc.page_content for doc in docs]
        metadatas = [doc.metadata for doc in docs]
        embeddings = self.embeddingFunction.encode_queries(contents)
        await self.mapper.insertBatchToEmbedding(
            [
                UserKnowledge(
                    knowledgeBaseId=knowledgeHeader.knowledgeBaseId,
                    userId=knowledgeHeader.userId,
                    content=content,
                    metadata=metadata,
                    embedding=embedding.detach().cpu().numpy(),
                ).model_dump()
                for content, metadata, embedding in zip(
                    contents, metadatas, embeddings.get("dense")
                )
            ]
        )
        await self.mapper.insertByEntity(
            UserKnowledgeBaseFile(
                knowledgeBaseId=knowledgeHeader.knowledgeBaseId,
                fileName=file_name,
                alias=_uuid,
                type=metadatas[0].get("type"),
            )
        )
        if metadatas[0].get("type") == "image":
            return contents[0]

    @override
    async def search(self, embeddingSearchParamsDTO: EmbeddingSearchParamsDTO):
        embedding = (
            self.embeddingFunction.encode_queries([embeddingSearchParamsDTO.query])
            .get("dense")[0]
            .detach()
            .cpu()
            .numpy()
        )
        know = await self.mapper.search(
            query_embeddiing=embedding,
            knowledgeBaseId=embeddingSearchParamsDTO.knowledgeBaseId,
            top_n=embeddingSearchParamsDTO.top_n,
            type=embeddingSearchParamsDTO.type,
        )
        know = sorted(know, key=lambda x: x.get("distance"), reverse=True)
        result = []
        for data in know:
            entity = data.get("entity")
            e = {
                "content": entity.get("content"),
                "type": entity.get("metadata").get("type"),
            }
            if e["type"] == "image":
                e["url"] = entity.get("metadata").get("url")
            if embeddingSearchParamsDTO.returnId:
                e["id"] = entity.get("metadata").get("uuid")
            result.append(e)
        return result

    @override
    async def getImages(self, knowledgeBaseId: int, count: int):
        # TODO 在表结构优化后，简化查询逻辑
        know = await self.mapper.getImagesInfo(knowledgeBaseId, count)
        result = []
        for data in know:
            metadata = data.get("metadata")
            result.append(
                ImageInfoVO(
                    url=metadata.get("url"),
                    id=metadata.get("uuid"),
                    fileName=metadata.get("file_name"),
                    content=data.get("content"),
                )
            )
        return result

    @override
    async def updateImageInfo(self, imageInfo: UpdateImageInfoDTO):
        exp = f"knowledgeBaseId == {imageInfo.knowledgeBaseId} and metadata['uuid'] == '{imageInfo.id}'"
        entity: dict = (await self.mapper.queryByExpression(exp))[0]
        if not entity:
            return
        id = entity.pop("id")
        await self.mapper.deleteEmbeddingById(id)
        embedding = (
            self.embeddingFunction.encode_queries([imageInfo.content])
            .get("dense")[0]
            .detach()
            .cpu()
            .numpy()
        )
        await self.mapper.insertBatchToEmbedding(
            [
                UserKnowledge(
                    knowledgeBaseId=imageInfo.knowledgeBaseId,
                    userId=entity.get("userId"),
                    content=imageInfo.content,
                    metadata=entity.get("metadata"),
                    embedding=embedding,
                ).model_dump()
            ]
        )


UserKnowledgeBaseService = T.Annotated[
    UserKnowledgeBaseServiceImpl, Depends(UserKnowledgeBaseServiceImpl())
]
